from datetime import datetime

import numpy as np
import torch
import torch.nn as nn


class Config(object):
    def __init__(self, base_dir, embedding):
        _now = datetime.now()
        self.model_name = "TextRNN"
        self.train_path = base_dir + "/data/train.txt"
        self.dev_path = base_dir + "/data/dev.txt"
        self.test_path = base_dir + "/data/test.txt"
        self.class_list = [x.strip() for x in open(
            base_dir + "/data/class.txt").readlines()]
        self.vocab_path = base_dir + "/data/vocab.pkl"
        self.save_path = base_dir + "/saved_dict/" + _now.strftime("%Y%m%d_%H_%M_%S") + self.model_name + ".ckpt"
        self.log_path = base_dir + "/log/" + self.model_name + "/" + _now.strftime("%Y%m%d_%H_%M_%S")
        self.embedding_pretrained = torch.tensor(
            np.load(base_dir + "/data/" + embedding)["embeddings"].astype("float32")
        )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.dropout = 0.5
        self.require_improvement = 10000
        self.num_classes = len(self.class_list)
        self.n_vocab = 8
        self.n_epochs = 10
        self.batch_size = 32
        self.pad_size = 15
        self.learning_rate = 1e-3
        self.embed_size = self.embedding_pretrained.shape[-1]
        self.hidden_size = 128
        self.num_layers = 3


class Model(nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        self.embedding = nn.Embedding.from_pretrained(config.embedding_pretrained, freeze=False)
        self.gru = nn.GRU(input_size=config.embed_size, hidden_size=config.hidden_size,
                           num_layers=config.num_layers, batch_first=True,
                           dropout=config.dropout, bidirectional=True
                           )
        self.fc = nn.Linear(config.hidden_size * 2, config.num_classes)

    def forward(self, x):
        embedded = self.embedding(x)
        # [batch, seq_len, embed]
        gru_out, _ = self.gru(embedded)
        # (batch_size, seq_len, hidden_size * num_directions)
        out = self.fc(gru_out[:, -1, :])
        return out



